[Frontend] DMA-descriptor rework + masked-DMA for non-dividing shapes (#282)#285
Merged
YWHyuk merged 13 commits intoJul 7, 2026
Merged
Conversation
…ptor
Replace the affine.apply{indirect_access} symbol smuggle with an explicit
offset descriptor. convert_indirect_indexing returns the offset spad instead
of folding sympy.Symbol(out) into the index; emit_transfer carries it as a
togsim.transfer operand; decompose_transfer lifts that operand to a
memref.dma_start {indirect_offset = @spad_symbol} attribute (memref.dma_start
is a registered op and rejects an extra operand, but accepts the attribute);
lower_dma_to_gemmini reads the attribute and resolves the global for CONFIG4
(drops _find_indirect); build_skeleton adds the offset spad to the gather DMA
read_bufs so the offset-build -> gather dependency edge forms in the trace.
The index stays clean (base only). Validated on both paths: Spike functional
(computed-index gather + scatter allclose, pointwise/reduce regression 0) and
the C++ trace timing path end-to-end (allclose; gather togsim_dma read_bufs now
includes the offset spad).
Indirect addressing (scattered-DMA timing) in the new trace path is a separate
gap tracked in issue #284.
Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
…ride, compute-loop multi-dim offset Three refinements on top of the explicit offset descriptor: 1. Detect indirect access by an explicit symbol set instead of an "tmp" substring match. indirect_indexing now records str(index_var) in self.indirect_symbols; a _has_indirect(expr) helper tests the index free_symbols against it; the former "tmp"-string sites (store, get_dma_info, convert entry/indirect_dims/stride) use it. Removes the now-dead _find_indirect from lower_dma_to_gemmini. 2. Single indirect dim: pass the raw index spad and let the MVIN apply the gather stride per position (CONFIG4 offset_stride) instead of a vector_load/muli/vector_store round-trip that pre-multiplied the stride. emit_transfer carries offset_stride; decompose copies the attribute; lower programs CONFIG4 with it (was hardcoded to 1). 3. Multi indirect dim (e.g. x[ix, iy]): the offset is a sum of strided indices, which a single CONFIG4 channel cannot do, so the sum stays in the kernel -- but build it in the compute loop (chunked by compute_vec_size, not a single tile-wide vector) and store it to a DEDICATED offset spad so an index that is live elsewhere (x[ix, iy] + ix) is not clobbered. push_step separates the offset-build loop from the gather that reads it. Adds multi-dim gather and index-reuse cases to test_indirect_access. Validated: indirect/scatter/embedding + the two new multi-dim cases pass; add/matmul/softmax regression 0. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
…emmini memref.dma_start is a registered op with a fixed operand list, so it cannot carry the runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers directly to Gemmini. - New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather. - togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without memref.dma_start's variadic tag indices. - memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul lowering, erased by the merged lowering). No memref.* DMA ops remain. - decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its DRAM-grow half is what the future masked-DMA clamp replaces). - dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer. Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile + async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv (padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0 passes; wrong values, not a crash). Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
…wait Finishes the memref.dma_start removal for the trace/timing path. build_tog's TogBuilder now decodes togsim.transfer (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) and togsim.wait (tag,tag_idx,num_elements) instead of memref.dma_start/dma_wait: op dispatch, _dma_start_fields (positional decode + direction from dma_kind / dma_type), the tag-user base-address scan, and the wait operand reads. TOGDMANode.tile_stride keeps reading the dram_stride attr (unchanged semantics). build_skeleton's wait marker/dispatch move to togsim.wait so barriers still become togsim.memory_barrier. dep_analysis needs no change (the decoder keeps the same field keys). Trace/timing (TOGSim, timing_mode=1) validated: add 5/5, matmul 11/11 (async DMA + tag_idx barrier pairing). No memref.* DMA ops remain in either path. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
… not operand count len(operands) > 8 is ambiguous once the masked-DMA clamp adds more optional operands (low/high after offset). Gate on the `indirect` attr instead, so offset vs mask operands are disambiguated by explicit flags. Behavior-preserving: gather/add pass on both the functional (Spike) and trace (TOGSim) paths. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
…gsim.wait) The trace producer now reads togsim.wait (not memref.dma_wait); update the pass docstring and comments to match. No code change. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
…o CONFIG/2/3/4) Instead of packing the transfer params into four CONFIG asm instructions, build a 144-byte DMA descriptor as an 18xi64 memref.global (.data, non-constant so runtime fields can be written) and emit a single CONFIG_DESC that hands its address to the MVIN/MVOUT (which read the struct, matching the Spike-side torchsim_config_desc). This is the vehicle for the masked-DMA low/high clamp: dim_low/dim_high/fill are descriptor fields (currently defaulted, so behavior is unchanged), not new special registers. - _pack_desc encodes the fields into the struct's i64 slots (dim_size/low/high/mm_stride/ spad_stride/element_size/vlane/config_type/flags/indirect stride+esize). - descriptor globals are deduped by content; the indirect offset spad address is a runtime value, stored into slot 15 before the transfer (flags.bit0 = indirect). Requires the Spike torchsim_config_desc instruction (riscv-isa-sim). Validated functional (Spike) + trace (TOGSim): add, matmul (subtile+async), softmax, gather/scatter (indirect runtime store), constant_pad all pass. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
The togsim.transfer descriptor lowering requires the spike torchsim_config_desc instruction (riscv-isa-sim PR #5, tag v1.0.4). Track the new tag so build-from-source picks up the descriptor-capable spike. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
The masked-DMA clamp needs spike to read the descriptor dim_low/dim_high + desc_fill, and index_add needs the MVOUT accumulate flag. Pin v1.0.5. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
…-dividing tails Drop the index_expr / indirect tile-divisibility recompile so pad_vlane_tile's lane-aligned tiles are kept instead of shrunk to a divisor of the dim. A non-dividing tile's per-tile remainder is handled by a masked-DMA clamp: - _emit_clamp emits per tile axis a dynamic [low, high) as affine.min/max of that axis' loop iv (the tile base), so the last partial tile and the pad borders fall out per iteration; low = max(0, glo - base), high = min(tile, ghi - base). The pointwise path (_masked_bounds) derives [glo, ghi) per dim (MVOUT [0, output); MVIN of a padded input [pad, pad+input_extent)); the MLIR templates (def_dma_op, e.g. gemm) derive it from the DRAM extent so a matmul tile padding past M/N/K is filled with 0 instead of MAC-ing garbage. - emit_transfer carries the (low, high) index operands + masked_axes; lower_transfer_to_gemmini writes the runtime low/high into the descriptor dim_low/dim_high through an i32-view (byte layout unchanged) + masked flag + desc_fill. One desc_ptr helper builds every descriptor: a fully static one is a deduped i64 global, while any runtime-written field (masked dim_low/dim_high and/or the indirect address, stored as two i32 words) uses a unique i32-view global, so a single descriptor carries both a masked clamp and a gather/scatter, and indirect DMAs drop their divisibility recompile too. - The box-excluded fill is the consuming reduction identity (reduction_init: sum 0, prod 1, max -inf, min +inf; 0 otherwise -- matmul MAC and stores use 0). Log-sum- exp exception (softmax / log_softmax): a sum whose node came from exp reduces exp(input), so the PRIMARY input (indexed by a reduction itervar) is filled with -inf; broadcast operands (the subtracted max) keep their finite identity, else input - (-inf) = +inf -> exp = inf -> NaN. - index_add: instead of the compute gather-add-overwrite (which loses duplicate indices landing in the same non-dividing tile), the MVOUT accumulates out[idx] += val via an accumulate flag; the DMA processes positions sequentially so duplicates accumulate correctly regardless of tile size. acc_float selects float vs integer add (element_size alone is ambiguous). With no caller left, is_dim_dividable / adjust_tile_to_divisible and the TileConstraint.must_divide_dim machinery (trim_large_tail bias, adjust loop) are removed. Fixes elementwise, constant_pad (mobilenet ragged vlane holes), reductions (sum/max/min/prod/softmax/log_softmax/var/std/layernorm/mean), gather/scatter (index_add with duplicates) and matmul/addmm on non-dividing shapes. Requires spike with the descriptor dim_low/dim_high clamp, desc_fill and the MVOUT accumulate flag. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
…ne-grained)
Conv on a non-dividing reduction (K = in_ch * kh * kw, e.g. 3*7*7=147 > systolic
size) MAC-ed garbage from the tail of the K tile. Two fixes:
- def_dma_op now takes the per-dim clamp extent from the kernel's loop_extents
{loop_iv_name: loop_bound}, which each conv template declares once (the template's
ranges/itervars equivalent). The stride-match extent inference is unreliable for
conv because the gemm reads a REPACKED weight (O_C contiguous, stride 1) whose
strides do not match the declared node_layout, so the inferred extent is wrong
(e.g. O_C=64 read as 8); gemm (no loop_extents) keeps the inference. Only a dim
whose index is a single iv is clamped -- an im2col spatial axis (o_h + k_h) has no
single base and is left unclamped, which is correct because that input is
pre-padded. All four conv templates (base, mt, sb, sbs) declare loop_extents; note
the mt template fuses channel and kernel-width so its tile_k loops 0..I_C*K_W.
- dma_fine_grained now PROPAGATES the masked (and indirect / accumulate) attrs +
the low/high operands when it splits a tile into subtiles, and REMAPS each clamp
per subtile: a subtile at offset iv*sub sees high = min(sub, high - iv*sub),
low = max(0, low - iv*sub) (out-of-window subtiles get a negative high / low > sub
and Spike skips them). Previously it dropped the clamp, so the conv clamp never
reached Spike; gemm was unaffected because its DMAs do not go through this pass.
Fixes conv2d on non-dividing in_channels / kernels (test_conv2d 13/13), including
fused pointwise epilogues (conv + relu/scale/bias/sigmoid). No regression:
matmul/addmm, bmm, elementwise, reduce, softmax, index_add all pass.
Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
A fused template epilogue (matmul/conv + relu/add/gelu/...) stores its result via store_epilogue, which -- unlike the plain store() -- emitted no masked-DMA clamp. On a non-dividing output the store then wrote the systolic pad region past the real extent, corrupting the result (fused matmul+relu at 10x10x10 was off by ~7). store_epilogue now clamps each output tile dim to [0, extent): the extent of every loop iv is ranges[k], and the tile dims are indexed by dim_aliasing (tile-dim order), so a dim whose iv extent does not divide the tile is clamped. Keyed by iv name, it is naming-agnostic -- it covers both the index-N matmul epilogue and the c0/tile_n/... conv epilogue -- and _emit_clamp skips dividing dims, so a dividing output is a no-op. Fixes fused matmul/conv epilogues on non-dividing shapes (matmul+relu/add/gelu, conv+relu/scale/bias/sigmoid, non-dividing K and non-dividing output). No regression: test_matmul, matmul_activation, addmm_residual, matmul_scalar all pass. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
Hand-write a masked MVOUT togsim.transfer, run lower_transfer_to_gemmini in-process (no torch, no spike; skipped without the MLIR bindings), and compare to an embedded golden -- the exact lowered IR (MLIR SSA numbering is deterministic for a fixed input): the i32-view descriptor global (dim_size / config_type / masked flag packed), the runtime dim_low/dim_high stores at i32 idx 6/10 for masked_axes=[2], and the CONFIG_DESC + MVOUT asm. A second case adds subtile_size=[1,1,4,8] and shows the descriptor dim_size become the subtile (config) shape. Timing mode erases the transfer. Under tests/lowering/ so further per-pass goldens have a home. Regenerate the golden and review the diff on an intended lowering change. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN
e6633a9 to
5c90140
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Base =
feature/codegen-ordered-steps.This PR started as "carry the gather/scatter offset as an explicit descriptor" (#282) and grew into the full DMA-descriptor rework on the
togsim.transferpath, plus masked-DMA support for non-dividing tile shapes.1. Gather/scatter offset as an explicit descriptor
convert_indirect_indexingreturns the offset spad; the DRAM index stays base-only (drops the+ Symbol(out)fold and the scalar-load/index_cast).emit_transfercarries the offset as atogsim.transferoperand, marked by anindirectattribute (detected by the attr, not the operand count).build_skeletonadds the offset spad to the gather DMA'sread_bufs, so the offset-build -> gather dependency edge forms in the trace.2. Lower
togsim.transferstraight to Gemmini + a DMA descriptormemref.dma_start:lower_transfer_to_gemminilowerstogsim.transferdirectly to Gemmini instructions (both the direct path and the trace/build_togpath, which is ported totogsim.transfer/togsim.wait).CONFIG*special-register writes with one TMA-style DMA descriptor: aCONFIG_DESCinstruction hands the descriptor pointer to the DMA engine, then MVIN/MVOUT reference it. Requires spike v1.0.4.3. Masked-DMA dynamic clamp (drop the tile-divisibility constraint)
is_dim_dividable/adjust_tile_to_divisible/TileConstraint.must_divide_dim. A non-dividing tile's tail is no longer avoided by shrinking the tile; instead each tile axis gets a per-dim[low, high)clamp written into the descriptordim_low/dim_highat run time (affine.min/affine.maxof the loop iv), and out-of-box positions are filled with the consuming reduction identity (sum 0, max -inf, ... ; matmul MAC / stores 0; log-sum-exp special-cased).index_addaccumulates via an MVOUT accumulate flag, so duplicate indices landing in the same tile add correctly.Validation (spike functional, regression 0)
tests/lowering/test_lower_transfer_to_gemmini.pycompares the masked-DMA descriptor lowering to an embedded golden.Follow-up
tests/ops/fusion/test_conv_fusion.pyfails with a spike subprocess crash at baseline too -- a separate, pre-existing issue, not this PR.🤖 Generated with Claude Code